#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import requests
import json
import logging
from typing import Dict, List, Optional
import os
from datetime import datetime

class DifyIntegration:
    """Dify LLM集成模块"""
    
    def __init__(self, api_key: str = None, api_base: str = None):
        self.logger = logging.getLogger(__name__)
        self.api_key = api_key or os.getenv('DIFY_API_KEY')
        self.api_base = api_base or os.getenv('DIFY_API_BASE')
        
        if not self.api_key or not self.api_base:
            raise ValueError("Dify API key and base URL must be provided")
            
        self.headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }
        
    def analyze_protocol(self, 
                        features: Dict,
                        behavior: Dict,
                        protocol_scores: Dict,
                        raw_packet: Optional[Dict] = None) -> Dict:
        """
        使用LLM分析VPN协议特征
        """
        # 构建提示词
        prompt = self._build_analysis_prompt(
            features, behavior, protocol_scores, raw_packet
        )
        
        try:
            # 调用Dify API
            response = requests.post(
                f"{self.api_base}/chat-messages",
                headers=self.headers,
                json={
                    "inputs": {},
                    "query": prompt,
                    "response_mode": "blocking",
                    "conversation_id": "",
                    "user": "vpn_analyzer"
                }
            )
            
            if response.status_code != 200:
                self.logger.error(f"Dify API error: {response.text}")
                return protocol_scores
                
            # 解析响应
            result = response.json()
            analysis = self._parse_llm_response(result)
            
            # 合并分析结果
            return self._merge_analysis(protocol_scores, analysis)
            
        except Exception as e:
            self.logger.error(f"Error calling Dify API: {str(e)}")
            return protocol_scores
            
    def _build_analysis_prompt(self,
                             features: Dict,
                             behavior: Dict,
                             protocol_scores: Dict,
                             raw_packet: Optional[Dict]) -> str:
        """构建分析提示词"""
        prompt = f"""作为网络安全专家，请分析以下VPN流量特征并给出协议识别建议：

1. 基本特征：
{json.dumps(features, indent=2)}

2. 行为特征：
{json.dumps(behavior, indent=2)}

3. 当前协议得分：
{json.dumps(protocol_scores, indent=2)}

"""

        if raw_packet:
            prompt += f"\n4. 原始数据包信息：\n{json.dumps(raw_packet, indent=2)}\n"

        prompt += """
请基于以上信息：
1. 分析最可能的VPN协议类型
2. 给出置信度评分（0-1）
3. 解释判断依据
4. 指出可疑特征

请以JSON格式返回，包含以下字段：
{
    "protocol": "最可能的协议",
    "confidence": 置信度,
    "reasoning": "判断依据",
    "suspicious_features": ["可疑特征列表"]
}
"""
        return prompt
        
    def _parse_llm_response(self, response: Dict) -> Dict:
        """解析LLM响应"""
        try:
            # 提取LLM的回复内容
            content = response.get('answer', '')
            
            # 尝试解析JSON
            if isinstance(content, str):
                # 查找JSON部分
                start = content.find('{')
                end = content.rfind('}') + 1
                if start >= 0 and end > start:
                    json_str = content[start:end]
                    return json.loads(json_str)
                    
            return {
                'protocol': 'unknown',
                'confidence': 0.0,
                'reasoning': 'Failed to parse LLM response',
                'suspicious_features': []
            }
            
        except Exception as e:
            self.logger.error(f"Error parsing LLM response: {str(e)}")
            return {
                'protocol': 'unknown',
                'confidence': 0.0,
                'reasoning': f'Error: {str(e)}',
                'suspicious_features': []
            }
            
    def _merge_analysis(self, 
                       protocol_scores: Dict,
                       llm_analysis: Dict) -> Dict:
        """合并LLM分析结果"""
        # 创建新的得分字典
        merged_scores = protocol_scores.copy()
        
        # 获取LLM识别的协议
        llm_protocol = llm_analysis.get('protocol', 'unknown')
        llm_confidence = llm_analysis.get('confidence', 0.0)
        
        # 如果LLM给出了明确的协议判断
        if llm_protocol != 'unknown' and llm_confidence > 0.5:
            # 调整得分
            for protocol in merged_scores:
                if protocol == llm_protocol:
                    # 提高LLM识别协议的得分
                    merged_scores[protocol] = min(
                        1.0,
                        merged_scores[protocol] + (llm_confidence * 0.3)
                    )
                else:
                    # 降低其他协议的得分
                    merged_scores[protocol] = max(
                        0.0,
                        merged_scores[protocol] - (llm_confidence * 0.1)
                    )
                    
        return merged_scores 